#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import torch
import logging

# 设置日志级别为DEBUG以查看详细信息
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 先记录初始的CUDA_VISIBLE_DEVICES
logger.info(f"初始 CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', '未设置')}")
logger.info(f"初始可用GPU数量: {torch.cuda.device_count() if torch.cuda.is_available() else 0}")

# 导入配置和ModelHandler
from config import Config
from model_handler import ModelHandler

# 确保使用GPU
Config.USE_GPU = True
Config.AUTO_SELECT_GPU = True
Config.GPU_COUNT = 1

# 1. 初始化ModelHandler并记录所有步骤
logger.info("\n===== 初始化ModelHandler =====")
model_handler = ModelHandler(Config)

# 打印GPU相关信息
logger.info(f"ModelHandler初始化后可用GPU数量: {torch.cuda.device_count()}")
logger.info(f"ModelHandler初始化后 CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', '未设置')}")
logger.info(f"ModelHandler选择的GPU: {model_handler.selected_gpus}")
logger.info(f"ModelHandler设备类型: {model_handler.device}")
logger.info(f"全局配置中的SELECTED_GPUS: {getattr(Config, 'SELECTED_GPUS', '未设置')}")

# 2. 尝试直接加载一个小模型进行测试
logger.info("\n===== 测试直接加载小模型 =====")
try:
    from transformers import AutoTokenizer, AutoModelForCausalLM
    
    # 使用一个非常小的测试模型
    test_model_path = "facebook/opt-125m"
    logger.info(f"尝试加载测试模型: {test_model_path}")
    logger.info(f"加载模型前 CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', '未设置')}")
    
    # 加载模型，设置device_map="auto"
    model = AutoModelForCausalLM.from_pretrained(
        test_model_path,
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True,
        device_map="auto"
    )
    
    # 查看模型实际加载到的设备
    model_device = next(model.parameters()).device
    logger.info(f"测试模型实际加载到的设备: {model_device}")
    
    # 清理内存
    del model
    torch.cuda.empty_cache()
except Exception as e:
    logger.error(f"加载测试模型失败: {str(e)}")

# 3. 测试原始ModelHandler的模型加载逻辑
logger.info("\n===== 测试ModelHandler的模型加载逻辑 =====")
try:
    # 尝试加载配置中的模型
    model_handler.load_model(Config.MODEL_PATH)
    
    # 查看模型实际加载到的设备
    if model_handler.model is not None:
        model_device = next(model_handler.model.parameters()).device
        logger.info(f"ModelHandler模型实际加载到的设备: {model_device}")
        
        # 检查model_handler.device是否与模型实际加载的设备一致
        logger.info(f"model_handler.device: {model_handler.device}")
        
        # 测试模型生成时使用的设备
        test_input = model_handler.tokenizer("你好", return_tensors="pt").to(model_handler.device)
        logger.info(f"测试输入张量的设备: {test_input.input_ids.device}")
    
except Exception as e:
    logger.error(f"ModelHandler加载模型失败: {str(e)}")

# 4. 测试环境变量设置对设备可见性的影响
logger.info("\n===== 测试环境变量设置对设备可见性的影响 =====")
for gpu_id in range(8):  # 假设系统有8个GPU
    try:
        # 临时设置CUDA_VISIBLE_DEVICES
        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
        logger.info(f"临时设置CUDA_VISIBLE_DEVICES={gpu_id}")
        
        # 检查torch能看到的设备数量
        torch.cuda.empty_cache()
        visible_count = torch.cuda.device_count()
        logger.info(f"临时设置后可见GPU数量: {visible_count}")
        
        if visible_count > 0:
            # 尝试创建一个张量并查看其设备
            tensor = torch.tensor([1.0], device='cuda')
            logger.info(f"在CUDA_VISIBLE_DEVICES={gpu_id}时创建的张量设备: {tensor.device}")
            
            # 获取当前设备的属性
            current_device = torch.cuda.current_device()
            prop = torch.cuda.get_device_properties(current_device)
            logger.info(f"当前设备属性: {prop.name}, 总内存: {prop.total_memory/1024**3:.2f}GB")
            
            # 清理内存
            del tensor
        
        # 恢复原始设置
        os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, model_handler.selected_gpus))
    except Exception as e:
        logger.error(f"测试GPU {gpu_id}时出错: {str(e)}")
        # 无论如何都恢复原始设置
        os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, model_handler.selected_gpus))

logger.info("\n调试完成！")